import SimpleITK as sitk
import numpy as np
from . import geoutils, fileutils
import math

"""
issues:
    1. 角度差或许可以直接算，而非考虑投影角度（旋转角度）    PS：旋转角度表达能力更强
    2. 距离差与角度差相关
"""

__all__ = ['calc_diff']


def nrrd2nifti(nrrd_path, nifti_path):
    image = sitk.ReadImage(nrrd_path)
    sitk.WriteImage(image, nifti_path)


def merge_image(path1, path2, out_path):
    image1 = sitk.ReadImage(path1)
    array1 = sitk.GetArrayFromImage(image1)
    image2 = sitk.ReadImage(path2)
    array2 = sitk.GetArrayFromImage(image2)

    array_out = array1.copy()
    for i in range(array1.shape[0]):
        for j in range(array1.shape[1]):
            for k in range(array1.shape[2]):
                if array2[i, j, k] != 0:
                    array_out[i, j, k] = 3.0

    image_out = sitk.GetImageFromArray(array_out)
    image_out.CopyInformation(image1)
    sitk.WriteImage(image_out, out_path)


def get_line_from_array(array: np.ndarray):
    """
    get analytical form (point, vector) from numerical form (array)
    NOTES: the XY projection shouldn't be too short. if so, swap axis first.
    maybe don't need so complicate

    :param array: 3D numerical array in order XYZ
    :return: point, vector in XYZ order
    """

    def get_center_y(x):
        y1, y2 = None, None
        for i in range(projection_xy.shape[1]):
            if projection_xy[x, i]:
                if y1 is None:
                    y1 = i
                y2 = i

        return (y1 + y2) / 2

    projection_xy = np.zeros((array.shape[0], array.shape[1]))
    for z in range(array.shape[2]):
        np.logical_or(projection_xy, array[:, :, z], out=projection_xy)
    x_start, x_end = 0, 0
    y_start, y_end = 0, 0
    x_bool = projection_xy.any(1)
    for i in range(projection_xy.shape[0]):
        if x_bool[i]:
            x_start = i
            break
    for i in range(projection_xy.shape[0] - 1, -1, -1):
        if x_bool[i]:
            x_end = i
            break
    y_bool = projection_xy.any(0)
    for i in range(projection_xy.shape[1]):
        if y_bool[i]:
            y_start = i
            break
    for i in range(projection_xy.shape[1] - 1, -1, -1):
        if y_bool[i]:
            y_end = i
            break

    # let X be the proper divide direction
    swapped = False
    if x_end - x_start < y_end - y_start:
        projection_xy = projection_xy.swapaxes(0, 1)
        swapped = True
    start = y_start if swapped else x_start
    end = y_end if swapped else x_end

    # divide the line-like area into pieces
    # the number should be odd to make the centerpoints be even
    piece_cnt = 11
    strip = (end - start) // piece_cnt
    projection_vector = np.zeros(2)
    sum_cnt = 0
    centerpoint_x = start + strip
    for i in range((piece_cnt - 1) // 2):
        # print(centerpoint_x, (x_start, x_end), (y_start, y_end))
        temp_center = centerpoint_x
        try:
            center = np.asarray([centerpoint_x, get_center_y(centerpoint_x)])
            centerpoint_x += strip
            center_end = np.asarray([centerpoint_x, get_center_y(centerpoint_x)])
            projection_vector += center_end - center
            sum_cnt += 1
            centerpoint_x += strip
        except:
            centerpoint_x = temp_center + 2 * strip

    projection_vector /= sum_cnt
    projection_vector /= np.linalg.norm(projection_vector)
    if swapped:
        projection_vector[0], projection_vector[1] = projection_vector[1], projection_vector[0]
    # the length may be not accurate!
    length = np.linalg.norm(np.asarray([x_start - x_end, y_start - y_end]))
    projection_vector *= length

    z_start, z_end = None, None
    if swapped:
        array = array.swapaxes(0, 1)

    z_bool = array[start].any(0)
    for i in range(array.shape[2]):
        if z_bool[i]:
            z_start = i
            break
    z_bool = array[end].any(0)
    for i in range(array.shape[2]):
        if z_bool[i]:
            z_end = i
            break

    vector = np.asarray([projection_vector[0], projection_vector[1], z_end - z_start])
    vector /= np.linalg.norm(vector)
    point = np.asarray([(x_start + x_end) / 2, (y_start + y_end) / 2, (z_start + z_end) / 2])

    return point, vector


def calc_diff(manual_axis_image_path, my_point=None, my_vector=None, my_axis_path=None, output_path=None):
    """
    calculate the difference between my axis and manual axis, write down some measurement value

    :param manual_axis_image_path: the path of manually labelled axis image.
    :param my_point: the point of my axis. if None, get it from .txt file
    :param my_vector: teh vector of my axis.
    :param my_axis_path: the path of .txt file which containing point \n vector inside
    :param output_path: the path to write down the output .txt file
    :return: the output measurement values: dist, (a0, b0), (a1, b1), theta
    """
    if my_point is None or my_vector is None:
        my_point, my_vector = fileutils.read_axis(my_axis_path)


    image = sitk.ReadImage(manual_axis_image_path)
    array = sitk.GetArrayFromImage(image)
    array = array.swapaxes(0, 2)
    manual_point, manual_vector = get_line_from_array(array)

    dist = geoutils.get_line_distance(my_point, my_vector, manual_point, manual_vector)
    a0, b0 = geoutils.get_rotation_angle(my_vector)
    a1, b1 = geoutils.get_rotation_angle(manual_vector)
    # not care the up/down direction first
    # reverse the manually labelled axis to make the theta be acute
    if math.acos(manual_vector.dot(my_vector)) > math.pi / 2:
        manual_vector *= -1
        a1, b1 = geoutils.get_rotation_angle(manual_vector)

    if output_path is not None:
        fileutils.write_metrics(output_path, my_point, my_vector, manual_point, manual_vector, dist, image.GetSpacing(), a0, b0, a1, b1)

    return dist, (a0, b0), (a1, b1), np.rad2deg(math.acos(manual_vector.dot(my_vector)))


if __name__ == '__main__':
    # TODO: test nrrd to nifti
    # nrrd_path = '00178-Segmentation.seg.nrrd'
    # nifti_path = '00178-Segmentation.seg.nii.gz'
    # nrrd2nifti(nrrd_path, nifti_path)

    # TODO: test if nrrd and nifti are the same --> yes, they are
    # array1 = sitk.GetArrayFromImage(sitk.ReadImage('Segmentation.seg.nrrd'))
    # array2 = sitk.GetArrayFromImage(sitk.ReadImage('Segmentation.seg.nii.gz'))
    # ans = np.logical_xor(array1, array2)
    # z = np.zeros(array1.shape)
    # print(np.all(ans == z))

    # TODO: test merge
    # image1_path = 'out_label.nii.gz'
    # mergeout_path = 'merged_label.nii.gz'
    # merge_image(image1_path, nifti_path, mergeout_path)

    # TODO: test get line and calc the diff
    # image = sitk.ReadImage('Segmentation.seg.nii.gz')
    # array = sitk.GetArrayFromImage(image)
    # array = array.swapaxes(0, 2)
    # point, vector = get_line_from_array(array)
    # print(point, vector)

    # # draw the line
    # array = array.swapaxes(0, 2)
    # step = 2
    # for i in range(-50, 51):
    #     temp_point = point + i * step * vector
    #     ind_i = round(temp_point[2])
    #     ind_j = round(temp_point[1])
    #     ind_k = round(temp_point[0])
    #     if ind_i >= array.shape[0] or ind_j >= array.shape[1] or ind_k >= array.shape[2]:
    #         break
    #     for j in range(-2, 3):
    #         for k in range(-2, 3):
    #             array[ind_i, ind_j + j, ind_k + k] = 2.0
    #
    # out_image = sitk.GetImageFromArray(array)
    # out_image.CopyInformation(image)
    # sitk.WriteImage(out_image, 'Segmentation_linefit.seg.nii.gz')


    # point_ = np.asarray([345.46317048, 9.9502225, 115.21751065])
    # vector_ = np.asarray([-0.21520773, 0.97033094, 0.11019755])
    # dist = get_line_distance(point_, vector_, point, vector)
    # a0, b0 = get_rotation_angle(vector_)
    # a1, b1 = get_rotation_angle(vector)
    # print('distance(pixel): ', dist)
    # print('spacing: ', image.GetSpacing())
    # print('distance(mm): ', dist * image.GetSpacing()[0])
    # print('angle of predicted: ', (np.rad2deg(a0), np.rad2deg(b0)))
    # print('angle of manual: ', (np.rad2deg(a1), np.rad2deg(b1)))
    # print('relative angle(degree): ', np.rad2deg(math.acos(vector.dot(vector_))))

    # TODO: test calc from file
    calc_diff('test_axis.txt', 'Segmentation.seg.nrrd', 'test_output.txt')